// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/memory/hash_string_allocator.h>
#include <pollux/exec/aggregate.h>
#include <pollux/expression/complex_view_types.h>
#include <pollux/expression/complex_writer_types.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/decoded_vector.h>

namespace kumo::pollux::aggregate {
    // Represents a list of values, including nulls, for an array/map/distinct value
    // set in aggregation. Bit-packed null flags are stored separately from the
    // non-null values.
    class ValueList {
    public:
        void appendValue(
            const DecodedVector &decoded,
            vector_size_t index,
            HashStringAllocator *allocator);

        template<typename T>
        void appendValue(
            const exec::OptionalAccessor<Generic<T> > &value,
            HashStringAllocator *allocator) {
            if (!value.has_value()) {
                appendNull(allocator);
            } else {
                POLLUX_DCHECK(!value->isNull());
                appendNonNull(*value->base(), value->decodedIndex(), allocator);
            }
        }

        void appendRange(
            const VectorPtr &vector,
            vector_size_t offset,
            vector_size_t size,
            HashStringAllocator *allocator);

        int32_t size() const {
            return size_;
        }

        // Called after 'finalize()' to get access to 'data' allocation.
        HashStringAllocator::Header *dataBegin() {
            return dataBegin_;
        }

        // Called after 'finalize()' to get access to 'nulls' allocation.
        HashStringAllocator::Header *nullsBegin() {
            return nullsBegin_;
        }

        uint64_t lastNulls() const {
            return lastNulls_;
        }

        void free(HashStringAllocator *allocator) {
            if (size_) {
                allocator->free(nullsBegin_);
                allocator->free(dataBegin_);
            }
        }

    private:
        // An array_agg or related begins with an allocation of 5 words and
        // 4 bytes for header. This is compact for small arrays (up to 5
        // bigints) and efficient if needs to be extended (stores 4 bigints
        // and a next pointer. This could be adaptive, with smaller initial
        // sizes for lots of small arrays.
        static constexpr int kInitialSize = 44;

        void appendNull(HashStringAllocator *allocator);

        void appendNonNull(
            const BaseVector &vector,
            vector_size_t index,
            HashStringAllocator *allocator);

        void prepareAppend(HashStringAllocator *allocator);

        // Writes lastNulls_ word to the 'nulls' block.
        void writeLastNulls(HashStringAllocator *allocator);

        // 'Nulls' allocation (potentially multi-part).
        HashStringAllocator::Header *nullsBegin_{nullptr};
        HashStringAllocator::Position nullsCurrent_{nullptr, nullptr};

        // 'Data' allocation (potentially multi-part)
        HashStringAllocator::Header *dataBegin_{nullptr};
        HashStringAllocator::Position dataCurrent_{nullptr, nullptr};

        // Number of values added, including nulls.
        uint32_t size_{0};

        // Bytes added. Used to control allocation of reserve for future appends.
        int32_t bytes_{0};

        // Last nulls word. 'size_ % 64' is the null bit for the next element.
        uint64_t lastNulls_{0};
    };

    // Extracts values from the ValueList into provided vector.
    class ValueListReader {
    public:
        explicit ValueListReader(ValueList &values);

        bool next(BaseVector &output, vector_size_t outputIndex);

    private:
        const vector_size_t size_;
        const vector_size_t lastNullsStart_;
        const uint64_t lastNulls_;
        HashStringAllocator::InputStream dataStream_;
        HashStringAllocator::InputStream nullsStream_;
        uint64_t nulls_;
        vector_size_t pos_{0};
    };

    // Write ValueList accumulators to Array-typed intermediate or final result
    // vectors.
    // TODO: This API only works if it is the only logic writing to `writer`.
    template<typename T>
    void copyValueListToArrayWriter(
        kumo::pollux::exec::ArrayWriter<T> &writer,
        ValueList &elements) {
        writer.resetLength();
        auto size = elements.size();
        if (size == 0) {
            return;
        }
        writer.reserve(size);

        ValueListReader reader(elements);
        for (vector_size_t i = 0; i < size; ++i) {
            reader.next(*writer.elementsVector(), writer.valuesOffset() + i);
        }
        writer.resize(size);
    }
} // namespace kumo::pollux::aggregate
